#ifndef _MS_H__
#define _MS_H__

//#define DEBUG   0
//#define S_DEBUG 0
/////////////////////////////////////////////////////
//
// MiaoSuan API
// Author: Zhang, Chen
// Tsinghua University
// 2019
//
/////////////////////////////////////////////////////

// Must call MS_init before any other MS API
void MS_init();

// Begin and end a MS session
// MS_begin();
// MS_end();
// These two APIs are macros and are placed in api.c

// Lock, Unlock
void MS_lock();

void MS_unlock();

/////////////////////////////////////////////////////
//
// Implementation ...
// Maybe there is a more elegant way ?
//
/////////////////////////////////////////////////////
#include <stdio.h>
#include <assert.h>
#include <string.h>
#include <slave.h>
#include "dma.h"
//#include "tprint.h"

//#define LOGE(...) tprintfci(TPRT_RED, __VA_ARGS__)
//#define LOGD(...) tprintfci(TPRT_BLUE, __VA_ARGS__)
//#define LOGDS(...) tprintfcx(TPRT_BLUE, 0, __VA_ARGS__)

///////////////////////////////////////////////
//
//  Basic Supports: Timing, DMA, Atomic, etc.
//
///////////////////////////////////////////////
#define FREQ (1425)

static inline uint64_t cpe_get_rtc() {
    unsigned long rpcc = 0;
    asm volatile ("rcsr %0, 4":"=&r"(rpcc)::"memory");
    return rpcc;
}

static inline uint64_t cpe_get_us() {
    unsigned long rpcc = 0;
    asm volatile ("rcsr %0, 4":"=&r"(rpcc)::"memory");
    return rpcc / FREQ;
}

static inline void cpe_dma_get(void *src, void *dst, size_t sz) {
    // src: global  dst: ldm
    if (sz == 0) return;
    volatile unsigned int reply = 0;
    athread_get(PE_MODE, src, dst, sz, (void *) &reply, 0, 0, 0);
    while (reply == 0);
}

static inline void cpe_dma_put(void *src, void *dst, size_t sz) {
    // src: ldm  dst: global
    if (sz == 0) return;
    volatile unsigned int reply = 0;
    athread_put(PE_MODE, src, dst, sz, (void *) &reply, 0, 0);
    while (reply == 0);
}

typedef volatile int spinlock_t;

static inline int cpe_atomic_inc(int *pos) {
    int rt;
    asm volatile ( "faaw %0, 0(%1) \n"
                   "memb"
    : "=&r" (rt)
    : "r" (pos));
    return rt;
}

inline void spin_lock(spinlock_t *lock_p) {
    unsigned int __tmp = 0;
    unsigned int __cnt;
    __asm__ __volatile__ (
    "0:     ldw     %[__tmp], %[lock_p]\n"
    "       beq     %[__tmp], 2f\n"
    "       ldi     %[__cnt], 1024\n"
    "       sll     %[__cnt], 4, %[__cnt]\n"
    "1:     subw    %[__cnt], 1, %[__cnt]\n"
    "       bne     %[__cnt], 1b\n"
    "       br      0b\n"
    "2:     faaw    %[__tmp], %[lock_p]\n"
    "       bne     %[__tmp], 0b\n"
    "       memb    \n"
    "       br      3f\n"
    "3:     unop    \n"
    : [__tmp] "=&r"(__tmp), [__cnt] "=&r"(__cnt)
    : [lock_p] "m"(*(lock_p))
    : "memory");
    return;
}

inline void spin_unlock(spinlock_t *lock_p) {
    unsigned int __tmp = 0;
    __asm__ __volatile__ (
    "       memb    \n"
    "       mov     0, %[__tmp]\n"
    "       stw     %[__tmp], %[lock_p]\n"
    : [__tmp] "=&r"(__tmp)
    : [lock_p] "m"(*(lock_p))
    : "memory"
    );
    return;
}

////////////////////////////////////////////////
//
//  Low Level Register Communication Support
//
////////////////////////////////////////////////
__thread_local_fix int msg_send_cnt, msg_recv_cnt;
#define MS_COMP_COL(x) (x & 0x07)
#define MS_COMP_ROW(x) ((x & 0x38) >> 3)
// #define REG_PUTR(var, dst) {asm volatile ("putr %0,%1\n"::"r"(var),"r"(dst)); msg_send_cnt ++;}
// #define REG_PUTC(var, dst) {asm volatile ("putc %0,%1\n"::"r"(var),"r"(dst)); msg_send_cnt ++;}
// #define REG_GETR(var) {asm volatile ("getr %0\n":"=r"(var)); msg_recv_cnt ++;}
// #define REG_GETC(var) {asm volatile ("getc %0\n":"=r"(var)); msg_recv_cnt ++;}

#define REG_PUTR(var, dst) {asm volatile ("putr %0,%1\n"::"r"(var),"r"(dst));}
#define REG_PUTC(var, dst) {asm volatile ("putc %0,%1\n"::"r"(var),"r"(dst));}
#define REG_GETR(var) {asm volatile ("getr %0\n":"=r"(var));}
#define REG_GETC(var) {asm volatile ("getc %0\n":"=r"(var));}

static inline int EMPTY_GETR(void) { int var; asm volatile ("rcsr %0,0x34\n":"=r"(var)); return var; }

static inline int EMPTY_GETC(void) { int var; asm volatile ("rcsr %0,0x35\n":"=r"(var)); return var; }

static inline int FULL_PUT(void) { int var; asm volatile ("rcsr %0,0x36\n":"=r"(var)); return var; }

////////////////////////////////////////////////
//
//  Message Encoding and Semantics
//  Helper data and functions
//
////////////////////////////////////////////////
__thread_local_fix int MS_TID, MS_ROW, MS_COL, MS_ROLE;
#define EXIT_FLAG 0x1
#define MSG_IS_EXIT(m) (m.msg.exit == EXIT_FLAG)
#define OWNER(x) (x&7)
#define OWNED(x) (OWNER(x) == MS_ROW)

typedef struct msg_t_ {
    uint8_t exit;
    uint8_t level; // 0 from client to server, 1 from server to server
    uint8_t src;
    uint8_t src_row;
    uint8_t src_col;
    uint8_t dst;
    uint8_t dst_row;
    uint8_t dst_col;
    // 64bit
    uint32_t lock_id;
    uint32_t operation;
    // 128bit
    uint64_t bitmap; // should use as uint8_t[8]
    uint64_t tail_padding;
} hmpmsg_t;
typedef union regmsg_t_ {
    uint256 v;
    hmpmsg_t msg;
} regmsg_t;
enum MS_OP {
    MS_INVALID = 0,
    MS_LOCK_ACQUIRE,
    MS_LOCK_RELEASE,
    MS_LOCK_APPROVE,
};
#define LOCK_MAX_WAIT_DEP  3
#define LOCK_REMOTE    0xC
#define LOCK_AVAILABLE 0x0
#define LOCK_OCCUPIED  0xF
#define LOCK_MULTI_MAX 64
volatile spinlock_t MS_MEM_LOCK[LOCK_MULTI_MAX];
//__thread_local_fix volatile regmsg_t ldm_buf[1024] __attribute((__aligned__(256))); // 32KB
__thread_local_fix volatile spinlock_t
MS_LDM_LOCK[LOCK_MULTI_MAX]; //__attribute((__aligned__(256)));
__thread_local_fix uint8_t
lock_chain[LOCK_MULTI_MAX<<LOCK_MAX_WAIT_DEP];
__thread_local_fix uint8_t
server_chain[LOCK_MULTI_MAX<<LOCK_MAX_WAIT_DEP];
__thread_local_fix uint8_t
lock_chain_tail[LOCK_MULTI_MAX], server_chain_tail[LOCK_MULTI_MAX];

////////////////////////////////////////////////
//
//  Server Client Mode Support
//  Auto Switch Support
//
////////////////////////////////////////////////
#define ROLE_CLIENT 0
#define ROLE_NOUSE  1
#define ROLE_SERVER 2
#define ROLE_META_SERVER 3

void MS_sc_enter() {
    msg_send_cnt = msg_recv_cnt = 0;
    if (MS_COL == 0) MS_ROLE = ROLE_SERVER;
    else MS_ROLE = ROLE_CLIENT;
    athread_syn(ARRAY_SCOPE, 0xffff);
//    tprintfm("My role is %2d\n", MS_ROLE);
}

void MS_sc_exit() {
    if (MS_ROLE == ROLE_CLIENT) {
        volatile regmsg_t m;
        m.msg.exit = EXIT_FLAG;
        REG_PUTR(m.v, 0);
//       tprintfi("[%d %d] client exit\n", MS_ROW, MS_COL);
    }
    /*
    tprintfm("%6d\n", msg_send_cnt);
    tprintfz("--------------------\n");
    tprintfm("%6d\n", msg_recv_cnt);
    tprintfz("--------------------\n");
    */
    athread_syn(ARRAY_SCOPE, 0xffff);
}

#define MSG_TEMPLATE(m, lock_id)\
  m.msg.exit = EXIT_FLAG - 1;\
  m.msg.src = MS_TID;\
  m.msg.src_row = MS_ROW;\
  m.msg.src_col = MS_COL;\
  m.msg.lock_id = lock_id;\



void do_server() {
    int exit_client = 0, exit_server = 0;
    int i;
    memset(server_chain, 0u, sizeof(server_chain));
    uint64_t row_id = MS_ROW;
    uint64_t server_chain_init = (row_id << 56) | (row_id << 48) | (row_id << 40) | (row_id << 32) | (row_id << 24) | (row_id << 16) | (row_id << 8) | row_id;
    for (i = MS_ROW; i < LOCK_MULTI_MAX; i += 8)
        ((uint64_t*)server_chain)[i] = server_chain_init;
    for (i = 0; i < LOCK_MULTI_MAX; i++)
        server_chain_tail[i] = OWNER(i);
    memset(lock_chain, 0u, sizeof(lock_chain));
    memset(lock_chain_tail, 0u, sizeof(lock_chain_tail));
    for (i = 0; i < LOCK_MULTI_MAX; i++)
        MS_LDM_LOCK[i] = LOCK_REMOTE;
    for (i = MS_ROW; i < LOCK_MULTI_MAX; i += 8)
        MS_LDM_LOCK[i] = LOCK_AVAILABLE;

    volatile regmsg_t m;
    while (exit_server < 8) {
        if (!EMPTY_GETC()) {
            REG_GETC(m.v);
            if (MSG_IS_EXIT(m)) exit_server++;
            else {
                uint32_t lock_id = m.msg.lock_id;
                if (OWNED(lock_id)) { // as metaserver
                    int op = m.msg.operation;
                    uint8_t client = m.msg.src_row;
                    if (op == MS_LOCK_ACQUIRE) {
                        if (MS_LDM_LOCK[lock_id] == LOCK_AVAILABLE) {
//                            tprintfic(TPRT_RED, "[%d-%d][Meta %d] give to server %d, %016llx\n", lock_id, OWNER(lock_id), MS_ROW, (int)client, ((uint64_t *) server_chain)[lock_id]);
                            MS_LDM_LOCK[lock_id] = LOCK_OCCUPIED;
                            m.msg.bitmap = ((uint64_t *) server_chain)[lock_id];
                            REG_PUTC(m.v, client);
                        } else {
                            server_chain[lock_id << LOCK_MAX_WAIT_DEP | server_chain_tail[lock_id]] = client;
                            server_chain_tail[lock_id] = client;
//                            tprintfic(TPRT_RED, "[%d-%d][Meta %d] recv req server %d, %016llx\n", lock_id, OWNER(lock_id), MS_ROW, (int)client, ((uint64_t *) server_chain)[lock_id]);
                        }
                    } else if (op == MS_LOCK_RELEASE) {
                        if (lock_chain_tail[lock_id] != 0) {
//                           tprintfic(TPRT_CYAN, "[%d-%d][Meta&Server %d] col release to client %d ,%016llx\n", lock_id, OWNER(lock_id), MS_ROW, (int)lock_chain[lock_id << LOCK_MAX_WAIT_DEP], ((uint64_t *) lock_chain)[lock_id]);
                            m.msg.bitmap = ((uint64_t *) lock_chain)[lock_id];
                            REG_PUTR(m.v, lock_chain[lock_id << LOCK_MAX_WAIT_DEP]);
                            lock_chain_tail[lock_id] = 0;
                            ((uint64_t *) lock_chain)[lock_id] = 0;
                        } else if (server_chain_tail[lock_id] != MS_ROW) {
//                            tprintfic(TPRT_RED, "[%d-%d][Meta %d] col release to server %d ,%016llx\n", lock_id, OWNER(lock_id), MS_ROW, (int)server_chain[lock_id << LOCK_MAX_WAIT_DEP|MS_ROW], ((uint64_t *) server_chain)[lock_id]);
                            m.msg.bitmap = ((uint64_t *) server_chain)[lock_id];
                            REG_PUTC(m.v, server_chain[lock_id << LOCK_MAX_WAIT_DEP|MS_ROW]);
                            server_chain_tail[lock_id] = MS_ROW;
                            ((uint64_t *) server_chain)[lock_id] = server_chain_init;
                        } else {
//                            tprintfic(TPRT_RED, "[%d-%d][Meta %d] col release\n", lock_id, OWNER(lock_id), MS_ROW);
                            MS_LDM_LOCK[lock_id] = LOCK_AVAILABLE;
                        }
                    }
                } else { // as server, is release
                    ((uint64_t *) server_chain)[lock_id] = m.msg.bitmap;
//                    tprintfic(TPRT_YELLOW, "[%d-%d][Server %d] Get from remote, server %016llx, lock %016llx\n", lock_id, OWNER(lock_id), MS_ROW, ((uint64_t *) server_chain)[lock_id], ((uint64_t *) lock_chain)[lock_id]);
                    m.msg.bitmap = ((uint64_t *) lock_chain)[lock_id];
                    REG_PUTR(m.v, lock_chain[lock_id << LOCK_MAX_WAIT_DEP]);
                    lock_chain_tail[lock_id] = 0;
                    ((uint64_t*) lock_chain)[lock_id] = 0;
                }
            }
        }
        if (!EMPTY_GETR()) {
            REG_GETR(m.v);
            if (MSG_IS_EXIT(m)) {
                exit_client++;
                if (exit_client == 7) {
                    exit_server++;
                    volatile regmsg_t m;
                    m.msg.exit = EXIT_FLAG;
//                    tprintfic(TPRT_RED, "[Server %d] all client exit\n", MS_ROW);
                    REG_PUTC(m.v, 8); // broadcast
                }
            } else {
                int op = m.msg.operation;
                uint8_t client = m.msg.src_col;
                uint32_t lock_id = m.msg.lock_id;
                if (OWNED(lock_id)) { // as server & metaserver
                    if (op == MS_LOCK_ACQUIRE) {
                        if (MS_LDM_LOCK[lock_id] == LOCK_OCCUPIED) {
                            lock_chain[lock_id << LOCK_MAX_WAIT_DEP | lock_chain_tail[lock_id]] = client;
                            lock_chain_tail[lock_id] = client;
//                            tprintfic(TPRT_CYAN, "[%d-%d][Meta&Server %d] recv req client %d, %016llx\n", lock_id, OWNER(lock_id), MS_ROW, (int)client, ((uint64_t *) lock_chain)[lock_id]);
                        } else { // available
                            m.msg.bitmap = ((uint64_t *) lock_chain)[lock_id];
//                           assert(m.msg.bitmap == 0);
//                            tprintfic(TPRT_CYAN, "[%d-%d][Meta&Server %d] give to client %d, %016llx\n", lock_id, OWNER(lock_id), MS_ROW, (int)client, ((uint64_t *) lock_chain)[lock_id]);
                            REG_PUTR(m.v, client);
                            MS_LDM_LOCK[lock_id] = LOCK_OCCUPIED;
                        }
                    } else { //release
                        if (server_chain_tail[lock_id] != MS_ROW) {
//                            tprintfic(TPRT_RED, "[%d-%d][Meta %d] row release to server %d ,%016llx\n", lock_id, OWNER(lock_id), MS_ROW, (int)server_chain[lock_id << LOCK_MAX_WAIT_DEP|MS_ROW], ((uint64_t *) server_chain)[lock_id]);
                            m.msg.bitmap = ((uint64_t *) server_chain)[lock_id];
                            REG_PUTC(m.v, server_chain[lock_id << LOCK_MAX_WAIT_DEP | MS_ROW]);
                            server_chain_tail[lock_id] = MS_ROW;
                            ((uint64_t *) server_chain)[lock_id] = server_chain_init;
                        } else if (lock_chain_tail[lock_id] != 0) {
//                            tprintfic(TPRT_CYAN, "[%d-%d][Meta&Server %d] row release to client %d ,%016llx\n", lock_id, OWNER(lock_id), MS_ROW, (int)lock_chain[lock_id << LOCK_MAX_WAIT_DEP], ((uint64_t *) lock_chain)[lock_id]);
                            m.msg.bitmap = ((uint64_t *) lock_chain)[lock_id];
                            REG_PUTR(m.v, lock_chain[lock_id << LOCK_MAX_WAIT_DEP]);
                            lock_chain_tail[lock_id] = 0;
                            ((uint64_t *) lock_chain)[lock_id] = 0;
                        } else  {
//                            tprintfic(TPRT_CYAN, "[%d-%d][Meta&Server %d] row release\n", lock_id, OWNER(lock_id), MS_ROW);
                            MS_LDM_LOCK[lock_id] = LOCK_AVAILABLE;
                        }
                    }
                } else { // not own the lock, do as server
                    if (op == MS_LOCK_ACQUIRE) {
                        if (MS_LDM_LOCK[lock_id] == LOCK_REMOTE) {
//                            tprintfic(TPRT_GREEN, "[%d-%d][Server %d] first lock, remote %d\n", lock_id, OWNER(lock_id), MS_ROW, (int)client);
                            REG_PUTC(m.v, OWNER(lock_id));
                            MS_LDM_LOCK[lock_id] = LOCK_OCCUPIED;
                        }
                        lock_chain[lock_id << LOCK_MAX_WAIT_DEP | lock_chain_tail[lock_id]] = client;
                        lock_chain_tail[lock_id] = client;
//                        tprintfic(TPRT_GREEN, "[%d-%d][Server %d] append lock, local %d, %016llx (%d)\n", lock_id, OWNER(lock_id), MS_ROW, (int)client, ((uint64_t *) lock_chain)[lock_id], (int)(lock_chain_tail[lock_id]));
                    } else { //release
//                        tprintfic(TPRT_GREEN, "[%d-%d][Server %d] pass to %d %016llx (%d)\n", lock_id, OWNER(lock_id), MS_ROW, server_chain[lock_id << LOCK_MAX_WAIT_DEP | MS_ROW], ((uint64_t *) server_chain)[lock_id], (int)server_chain_tail[lock_id]);
                        m.msg.bitmap = ((uint64_t *) server_chain)[lock_id];
                        REG_PUTC(m.v, server_chain[lock_id << LOCK_MAX_WAIT_DEP | MS_ROW]);
                        if (lock_chain_tail[lock_id] == 0) {
//                            tprintfic(TPRT_GREEN, "[%d-%d][Server %d] free\n", lock_id, OWNER(lock_id), MS_ROW);
                            MS_LDM_LOCK[lock_id] = LOCK_REMOTE;
                        } else {
//                            tprintfic(TPRT_GREEN, "[%d-%d][Server %d] and remote %016llx (%d)\n", lock_id, OWNER(lock_id), MS_ROW, ((uint64_t *) lock_chain)[lock_id], (int)lock_chain_tail[lock_id]);
                            m.msg.operation = MS_LOCK_ACQUIRE;
                            REG_PUTC(m.v, OWNER(lock_id));
                        }
                    }
                }
            }
        }
    }
//    tprintfic(TPRT_MAGENTA, "[Server %d] exit\n", MS_ROW);
}
#define MS_begin() \
    MS_sc_enter();\
    if (MS_ROLE == ROLE_CLIENT)\
    {

#define MS_end()\
    }\
    else if (MS_ROLE == ROLE_NOUSE) {}\
    else if (MS_ROLE == ROLE_SERVER) do_server();\
    MS_sc_exit();\


inline void MS_lock(uint32_t lock_id) {
//   tprintfi("[%d-%d][%d %d] apply lock\n", lock_id, OWNER(lock_id), MS_ROW, MS_COL);
    volatile regmsg_t m;
    MSG_TEMPLATE(m, lock_id);
    m.msg.operation = MS_LOCK_ACQUIRE;
    REG_PUTR(m.v, 0);
    REG_GETR(m.v);
    ((uint64_t *) lock_chain)[lock_id] = m.msg.bitmap;
//   tprintfi("[%d-%d][%d %d] get lock %016llx\n", lock_id, OWNER(lock_id), MS_ROW, MS_COL, ((uint64_t*)lock_chain)[lock_id]);
}

inline void MS_unlock(uint32_t lock_id) {
    volatile regmsg_t m;
    MSG_TEMPLATE(m, lock_id);
    m.msg.operation = MS_LOCK_RELEASE;
    m.msg.bitmap = ((uint64_t *) lock_chain)[lock_id];
    REG_PUTR(m.v, lock_chain[lock_id << LOCK_MAX_WAIT_DEP | MS_COL]);
//   tprintfi("[%d-%d][%d %d] release %016llx\n", lock_id, OWNER(lock_id), MS_ROW, MS_COL, ((uint64_t*)lock_chain)[lock_id]);
}

volatile int barrier_flag;

inline void MS_barrier() {
    int MS_CLIENT_ID = (MS_ROW ) * 7 + (MS_COL - 1);
    cpe_atomic_inc(&barrier_flag);
    while (barrier_flag % 56);
}

inline void MS_init(void) {
    barrier_flag = 0;
    MS_TID = athread_get_id(-1);
    MS_ROW = MS_COMP_ROW(MS_TID);
    MS_COL = MS_COMP_COL(MS_TID);
    int i;
    if (MS_COL == 0) {
        for (i = MS_ROW; i < LOCK_MULTI_MAX; i+=8)
            MS_MEM_LOCK[i] = LOCK_AVAILABLE;
    }
    if (MS_TID == 0)
        printf("Approach = Merged Lock\n");
}

#endif
